/*
 * Copyright 2013 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifdef __MINGW64__
#undef UNICODE
#undef _UNICODE
#endif
#ifdef DEBUG
#include <stdio.h>
#endif
#include <stddef.h>
#include <stdlib.h>
#include <tchar.h>
#define WIN32_LEAN_AND_MEAN 1
#include <windows.h>

#include "reimport.h"

#include "growing_buffer.h"
#include "rva.h"

const char* usage =
		"Usage: reimport [input.exe|input.dll] [output.exe|output.dll] hook.dll [library.dll:Function]...\n"
				"  redirect functions in input's import table to hook.dll\n";

typedef struct {
	HANDLE hFile;
	HANDLE hMapping;
	char* mapping;
	SIZE_T mapSize;
} mapped_file;

char* ConvertString(_TCHAR* str) {
#ifdef _UNICODE
	int len;
	char* result;
	len = wcslen(str);
	result = malloc(2 * len + 1);
	wcstombs(result, str, 2 * len + 1);
	return result;
#else
	return strdup(str);
#endif
}

int AlignUp(int size, int align) {
	return ((size + align - 1) / align) * align;
}

void PrintMessage(HANDLE file, const char* message) {
	int len = strlen(message);
	DWORD written;
	while (len > 0) {
		WriteConsoleA(file, message, len, &written, NULL);
		len -= written;
		message += written;
	}
}

void ErrorMessage(const char* message) {
	PrintMessage(GetStdHandle(STD_ERROR_HANDLE), message);
}

void OutputMessage(const char* message) {
	PrintMessage(GetStdHandle(STD_OUTPUT_HANDLE), message);
}

/*
 * Open file and map it to memory. Returns TRUE on success and FALSE on error.
 */
BOOL OpenMappedFile(mapped_file* in, _TCHAR* fileName) {
	MEMORY_BASIC_INFORMATION memInfo;
	in->hFile = NULL;
	in->hMapping = NULL;
	in->mapping = NULL;
	in->hFile = CreateFile(fileName, GENERIC_READ, FILE_SHARE_READ, NULL,
			OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
	if (in->hFile == INVALID_HANDLE_VALUE) {
		ErrorMessage("Can't open input file\n");
		return FALSE;
	}
	in->hMapping = CreateFileMapping(in->hFile, NULL, PAGE_READONLY, 0, 0,
			NULL);
	if (in->hMapping == NULL) {
		ErrorMessage("Can't create file mapping\n");
		return FALSE;
	}
	in->mapping = MapViewOfFile(in->hMapping, FILE_MAP_READ, 0, 0, 0);
	if (in->mapping == NULL) {
		ErrorMessage("Can't map file to memory\n");
		return FALSE;
	}
	if (VirtualQuery(in->mapping, &memInfo, sizeof(memInfo)) == 0) {
		ErrorMessage("Fail to get size of file mapping\n");
		return FALSE;
	}
	in->mapSize = memInfo.RegionSize;
	return TRUE;
}

/*
 * Close file mapped to memory. File mapping can be partially initialized.
 */
void CloseMappedFile(mapped_file* in) {
	if (in->mapping != NULL) {
		UnmapViewOfFile(in->mapping);
	}
	if (in->hMapping != NULL) {
		CloseHandle(in->hMapping);
	}
	if (in->hFile != NULL) {
		CloseHandle(in->hFile);
	}
}

typedef struct {
	/* we are working on 32 or 64 bit program. */
	int bits;
	/* PE header offset in the file. */
	int headerOffset;
	/* Number of sections in original executable. */
	int numberOfSections;
	/* Section alignment in file. */
	int fileAlignment;
	/* Section alignment in memory. */
	int sectionAlignment;
	/* Size of image in memory of original executable. */
	int sizeOfImage;
	/* Last position in original file that corresponds to some section. */
	int sizeOfImageInFile;
	/* Pointer to sections data in original executable. */
	PIMAGE_SECTION_HEADER sections;
	/* Start offset in file of sections data in original executable. */
	int sectionStartOffset;
	/* End offset in file of sections data in original executable. */
	int sectionEndOffset;
	/* Offset of sections data in file in new executable relative old executable. */
	int fileOffset;
	/* Offset of sections data in memory in new executable relative old executable. */
	int rvaOffset;
	/* Pointer to directory data in original executable. */
	IMAGE_DATA_DIRECTORY* directory;
	/* Pointer to original executable mapped to memory. */
	char* mapping;
	/* New section that contains new import tables. */
	growing_buffer newSection;
	/* Offset of new section in memory. */
	DWORD newSectionRVA;
	/* Offset of import tables in memory in new executable. */
	DWORD importTablesRVA;
	/* Number of import tables in original executable. */
	int numberOfImportTables;
	/* Number of import tables in new executable. */
	int newNumberOfImportTables;
	/* Point to import table in original executable. */
	PIMAGE_IMPORT_DESCRIPTOR importTables;
} pe_info;

/*
 * Calculate result->fileOffset and result->rvaOffset.
 */
void CalculateOffsets(pe_info* result, int sizeOfHeaders) {
	int newSize;
	int headers;
	int firstSectionRVA;
	newSize = result->sectionEndOffset + sizeof(IMAGE_SECTION_HEADER);
	headers = AlignUp(sizeOfHeaders, result->fileAlignment);
	newSize = AlignUp(newSize, result->fileAlignment);
	if (newSize > headers) {
		result->fileOffset = newSize - headers;
	}
	firstSectionRVA = AlignUp(headers, result->sectionAlignment);
	newSize = AlignUp(newSize, result->sectionAlignment);
	if (newSize > firstSectionRVA) {
		result->rvaOffset = newSize - firstSectionRVA;
	}
}

/*
 * Calculate result->sizeOfImageInFile.
 */
void CalculateSizeOfImageInFile(pe_info* result) {
	int i;
	int fileEnd;
	result->sizeOfImageInFile = result->sectionEndOffset;
	for (i = 0; i < result->numberOfSections; i++) {
		int fileEnd = result->sections[i].PointerToRawData
				+ AlignUp(result->sections[i].SizeOfRawData,
						result->fileAlignment);
		if (fileEnd > result->sizeOfImageInFile) {
			result->sizeOfImageInFile = fileEnd;
		}
#ifdef DEBUG
		printf("%x %x %x %x\n",
				(unsigned int) result->sections[i].VirtualAddress,
				(unsigned int) result->sections[i].Misc.VirtualSize,
				(unsigned int) result->sections[i].PointerToRawData,
				(unsigned int) result->sections[i].SizeOfRawData);
#endif
	}
}

/*
 * Calculate result->sectionStartOffset and result->sectionEndOffset.
 */
void CalculateSectionInfoPosition(pe_info* result) {
	result->sectionStartOffset = ((char*) result->sections
			- (char*) result->mapping);
	result->sectionEndOffset = result->sectionStartOffset
			+ result->numberOfSections * sizeof(IMAGE_SECTION_HEADER);
}

/*
 * Parse header of 32-bit input program and fill result fields.
 */
BOOL ParseInputFileHeader32(mapped_file* in, PIMAGE_NT_HEADERS32 peHeader,
		pe_info* result) {
	result->bits = 32;
	result->mapping = in->mapping;
	result->numberOfSections = peHeader->FileHeader.NumberOfSections;
	result->fileAlignment = peHeader->OptionalHeader.FileAlignment;
	result->sectionAlignment = peHeader->OptionalHeader.SectionAlignment;
	result->sizeOfImage = peHeader->OptionalHeader.SizeOfImage;
	result->sections = IMAGE_FIRST_SECTION(peHeader);
	result->fileOffset = 0;
	result->rvaOffset = 0;
	result->directory =
			(IMAGE_DATA_DIRECTORY*) &peHeader->OptionalHeader.DataDirectory;
	CalculateSectionInfoPosition(result);
	CalculateOffsets(result, peHeader->OptionalHeader.SizeOfHeaders);
	CalculateSizeOfImageInFile(result);
	return TRUE;
}

/*
 * Parse header of 64-bit input program and fill result fields.
 */
BOOL ParseInputFileHeader64(mapped_file* in, PIMAGE_NT_HEADERS64 peHeader,
		pe_info* result) {
	result->bits = 64;
	result->mapping = in->mapping;
	result->numberOfSections = peHeader->FileHeader.NumberOfSections;
	result->fileAlignment = peHeader->OptionalHeader.FileAlignment;
	result->sectionAlignment = peHeader->OptionalHeader.SectionAlignment;
	result->sizeOfImage = peHeader->OptionalHeader.SizeOfImage;
	result->sections = IMAGE_FIRST_SECTION(peHeader);
	result->fileOffset = 0;
	result->rvaOffset = 0;
	result->directory =
			(IMAGE_DATA_DIRECTORY*) &peHeader->OptionalHeader.DataDirectory;
	CalculateSectionInfoPosition(result);
	CalculateOffsets(result, peHeader->OptionalHeader.SizeOfHeaders);
	CalculateSizeOfImageInFile(result);
	return TRUE;
}

/*
 * Parse header of input program and fill result fields.
 */
BOOL ParseInputFileHeader(mapped_file* in, pe_info* result) {
	PIMAGE_DOS_HEADER dosHeader;
	PIMAGE_NT_HEADERS peHeader;
	dosHeader = (PIMAGE_DOS_HEADER) in->mapping;
	if (dosHeader->e_magic != IMAGE_DOS_SIGNATURE) {
		ErrorMessage("Invalid file format. DOS header not found.\n");
		return FALSE;
	}
	result->headerOffset = dosHeader->e_lfanew;
	peHeader = (PIMAGE_NT_HEADERS) (in->mapping + result->headerOffset);
	if (peHeader->Signature != IMAGE_NT_SIGNATURE) {
		ErrorMessage("Invalid file format. PE header not found.\n");
		return FALSE;
	}
	if (peHeader->OptionalHeader.Magic == IMAGE_NT_OPTIONAL_HDR32_MAGIC) {
		return ParseInputFileHeader32(in, (PIMAGE_NT_HEADERS32) peHeader,
				result);
	} else if (peHeader->OptionalHeader.Magic == IMAGE_NT_OPTIONAL_HDR64_MAGIC) {
		return ParseInputFileHeader64(in, (PIMAGE_NT_HEADERS64) peHeader,
				result);
	}
	ErrorMessage(
			"Invalid file format. Unknown optional header magic number.\n");
	return FALSE;
}

/*
 * Initialize pe_sections fields using pe_info structure.
 */
void InitPeSections(pe_sections* peSections, pe_info* peInfo) {
	peSections->sections = peInfo->sections;
	peSections->numberOfSections = peInfo->numberOfSections;
	peSections->mapping = peInfo->mapping;
	peSections->fileAlignment = peInfo->fileAlignment;
	peSections->sectionAlignment = peInfo->sectionAlignment;
}

/*
 * Copy IMAGE_IMPORT_DESCRIPTORs from old IAT to new IAT.
 */
BOOL CopyImportDescriptors(rva_pointer* pIAT, growing_buffer* newSection,
		pe_info* peInfo) {
	rva_pointer p;
	IMAGE_IMPORT_DESCRIPTOR importDescriptor;
	PIMAGE_IMPORT_DESCRIPTOR newSectionData;
	p = *pIAT;
	peInfo->importTables = GetBufferEnd(newSection);
	peInfo->numberOfImportTables = 0;
	while (1) {
		if (!ReadData(&p, &importDescriptor, sizeof(importDescriptor))) {
			break;
		}
		if (importDescriptor.OriginalFirstThunk == 0
				&& importDescriptor.FirstThunk == 0) {
			break;
		}
		peInfo->numberOfImportTables++;
		newSectionData = GrowBuffer(newSection, sizeof(importDescriptor));
		memcpy(newSectionData, &importDescriptor, sizeof(importDescriptor));
		if (!SetRVA(&p, p.rva + sizeof(importDescriptor))) {
			ErrorMessage("Invalid import table");
			return FALSE;
		}
	}
	return TRUE;
}

/*
 * Add zero IMAGE_IMPORT_DESCRIPTOR that marks the end of list of import tables.
 */
void AddTerminationImportDescriptor(growing_buffer* newSection) {
	PIMAGE_IMPORT_DESCRIPTOR emptyDescriptor;
	emptyDescriptor = GrowBuffer(newSection, sizeof(IMAGE_IMPORT_DESCRIPTOR));
	memset(emptyDescriptor, 0, sizeof(IMAGE_IMPORT_DESCRIPTOR));
}

/*
 * Store hook dll name in newSection. Return RVA of the string.
 */
DWORD AddHookDllName(pe_info* peInfo, growing_buffer* newSection,
		_TCHAR* hookDllName) {
	DWORD hookDllRVA;
	char* hookDll;
	size_t hookDllByteLength;
	char* hookDllInSection;
	hookDllRVA = newSection->size + peInfo->newSectionRVA;
	hookDll = ConvertString(hookDllName);
	hookDllByteLength = strlen(hookDll) + 1;
	hookDllInSection = GrowBuffer(&peInfo->newSection, hookDllByteLength);
	memcpy(hookDllInSection, hookDll, hookDllByteLength);
	free(hookDll);
	return hookDllRVA;
}

/*
 * Search function in an import table. Return TRUE if function is found,
 * RVA of its name in functionNameRVA and its index in functionIndex.
 * Return FALSE on errors or if function is not found.
 */
BOOL FindFunctionInLibrary(pe_sections* peSections, int OriginalFirstThunk,
		char* function, int* functionIndex, DWORD* functionNameRVA) {
	rva_pointer p;
	rva_pointer name;
	DWORD lookup;
	char* functionName;
	int found;
	InitRVAPointer(&p, peSections, OriginalFirstThunk);
	InitRVAPointer(&name, peSections, OriginalFirstThunk);
	*functionNameRVA = 0;
	for (*functionIndex = 0;; (*functionIndex)++) {
		if (!ReadDWORD(&p, &lookup)) {
			return FALSE;
		}
		if (!SetRVA(&p, p.rva + 4)) {
			ErrorMessage("Invalid original thunk.");
			return FALSE;
		}
		if ((lookup & 0x80000000u) != 0) {
			continue;
		}
		if (lookup == 0) {
			*functionIndex = -1;
			break;
		}
		/* First 2 bytes in function name are hint for linker */
		if (!SetRVA(&name, lookup + 2)) {
			ErrorMessage("Invalid original thunk value.");
			return FALSE;
		}
		functionName = ReadString(&name);
		if (functionName == NULL) {
			ErrorMessage("Invalid original thunk function name.");
			return FALSE;
		}
		found = strcmp(functionName, function);
		free(functionName);
		if (found == 0) {
			*functionNameRVA = lookup;
			break;
		}
	}
	return TRUE;
}

/*
 * Store the name of library function. Return RVA of stored string.
 */
DWORD AppendFunctionName(pe_info* peInfo, char* function,
		growing_buffer* newSection) {
	DWORD functionRVA;
	char* functionDst;
	if ((newSection->size & 1) != 0) {
		GrowBuffer(newSection, 1);
	}
	functionRVA = newSection->size + peInfo->newSectionRVA;
	/*
	 * We don't know hint for linker, so we set it to zero.
	 * Linker rechecks hint anyway.
	 */
	*(short*) GrowBuffer(newSection, sizeof(short)) = 0;
	functionDst = GrowBuffer(newSection, strlen(function) + 1);
	strcpy(functionDst, function);
	if ((newSection->size & 3) != 0) {
		GrowBuffer(newSection, 4 - (newSection->size & 3));
	}
	return functionRVA;
}

/*
 * Add new import table that redirects one import.
 * tableOffset - offset of import table inside new section. We use
 * offset instead of pointer, since new section can move in memory
 * when it grows.
 */
BOOL AddImportTable(pe_info* peInfo, pe_sections* peSections,
		growing_buffer* newSection, DWORD hookRVA, int tableOffset,
		_TCHAR* func) {
	int i;
	char* funcStr;
	char* library;
	char* function;
	char* newFunction;
	PIMAGE_IMPORT_DESCRIPTOR importTable;
	char* importTableLibraryName;
	rva_pointer p;
	int foundLibrary;
	int functionIndex;
	DWORD functionNameRVA;
	DWORD* originalThunk;
	PIMAGE_IMPORT_DESCRIPTOR table;

	table = (PIMAGE_IMPORT_DESCRIPTOR) GetBufferPointer(newSection,
			tableOffset);
	table->TimeDateStamp = -1;
	table->ForwarderChain = 0;
	table->Name = hookRVA;

	funcStr = ConvertString(func);
	function = strchr(funcStr, ':');
	if (function == NULL) {
		ErrorMessage(
				"Invalid function definition format. Should be library.dll:function");
		free(funcStr);
		return FALSE;
	}
	*function = 0;
	library = funcStr;
	function++; /* skip ":" */
	newFunction = strchr(function, ':');
	if (newFunction != NULL) {
		*newFunction = 0;
		newFunction++; /* skip ":" */
	}

	for (i = 0; i < peInfo->numberOfImportTables; i++) {
		importTable = &peInfo->importTables[i];
		InitRVAPointer(&p, peSections, importTable->Name);
		importTableLibraryName = ReadString(&p);
		if (importTableLibraryName == NULL) {
			ErrorMessage("Can't read import library name.");
			free(funcStr);
			return FALSE;
		}
		foundLibrary = stricmp(importTableLibraryName, library);
		free(importTableLibraryName);
		if (foundLibrary == 0) {
			if (!FindFunctionInLibrary(peSections,
					importTable->OriginalFirstThunk, function, &functionIndex,
					&functionNameRVA)) {
				free(funcStr);
				return FALSE;
			}
			if (functionIndex >= 0) {
				if (newFunction != NULL) {
					functionNameRVA = AppendFunctionName(peInfo, newFunction,
							newSection);
					/* recalculate pointer in case newSection was reallocated */
					table = (PIMAGE_IMPORT_DESCRIPTOR) GetBufferPointer(
							newSection, tableOffset);
				}

				table->FirstThunk = importTable->FirstThunk + functionIndex * 4;
				table->OriginalFirstThunk = newSection->size
						+ peInfo->newSectionRVA;
				originalThunk = GrowBuffer(newSection, sizeof(DWORD));
				*originalThunk = functionNameRVA;
				originalThunk = GrowBuffer(newSection, sizeof(DWORD));
				*originalThunk = 0;
				free(funcStr);
				return TRUE;
			}
		}
	}
	ErrorMessage("A function is not found in the import.");
	table->OriginalFirstThunk = newSection->size + peInfo->newSectionRVA;
	originalThunk = GrowBuffer(newSection, sizeof(DWORD));
	*originalThunk = 0;
	table->FirstThunk = table->OriginalFirstThunk; /* Can be anything. */
	free(funcStr);
	return TRUE;
}

BOOL AddImportTables(mapped_file* in, pe_info* peInfo, _TCHAR** argv, int argc) {
	pe_sections peSections;
	DWORD IAT;
	rva_pointer pIAT;
	DWORD hookRVA;
	int tablesOffset;
	int i;

	peInfo->newSectionRVA = peInfo->sizeOfImage + peInfo->rvaOffset;
	CreateBuffer(&peInfo->newSection);
	hookRVA = AddHookDllName(peInfo, &peInfo->newSection, argv[0]);
	InitPeSections(&peSections, peInfo);
	IAT = peInfo->directory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress;
	InitRVAPointer(&pIAT, &peSections, IAT);
	peInfo->importTablesRVA = peInfo->newSectionRVA + peInfo->newSection.size;
	if (!CopyImportDescriptors(&pIAT, &peInfo->newSection, peInfo)) {
		return FALSE;
	}
	tablesOffset = peInfo->newSection.size;
	GrowBuffer(&peInfo->newSection,
			(argc - 1) * sizeof(IMAGE_IMPORT_DESCRIPTOR));
	AddTerminationImportDescriptor(&peInfo->newSection);
	peInfo->newNumberOfImportTables = peInfo->numberOfImportTables + argc - 1;
	for (i = 1; i < argc; i++) {
		if (!AddImportTable(peInfo, &peSections, &peInfo->newSection, hookRVA,
				tablesOffset + (i - 1) * sizeof(IMAGE_IMPORT_DESCRIPTOR),
				argv[i])) {
			return FALSE;
		}
	}
	return TRUE;
}

BOOL CreateOutputFile(mapped_file* in, mapped_file* out, pe_info* peInfo,
		_TCHAR* outputFileName) {
	LARGE_INTEGER inputFileSize;
	LARGE_INTEGER outputFileSize;
	LARGE_INTEGER outputSectionsSize;
	DWORD bytesWritten;
	int sizeIncrease;
	if (!GetFileSizeEx(in->hFile, &inputFileSize)) {
		ErrorMessage("Can't determine input file size");
		return FALSE;
	}
	sizeIncrease = peInfo->fileOffset
			+ AlignUp(peInfo->newSection.size, peInfo->fileAlignment);
	out->hFile = CreateFile(outputFileName, GENERIC_READ | GENERIC_WRITE,
			FILE_SHARE_READ, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
	if (out->hFile == INVALID_HANDLE_VALUE) {
		return FALSE;
	}
	out->hMapping = NULL;
	out->mapping = NULL;
	outputFileSize.QuadPart = inputFileSize.QuadPart + sizeIncrease;
	if (!SetFilePointerEx(out->hFile, outputFileSize, NULL, FILE_BEGIN)) {
		CloseMappedFile(out);
		return FALSE;
	}
	if (!SetEndOfFile(out->hFile)) {
		CloseMappedFile(out);
		return FALSE;
	}
	if (outputFileSize.QuadPart > peInfo->sizeOfImageInFile) {
		outputSectionsSize.QuadPart = peInfo->sizeOfImageInFile + sizeIncrease;
		if (!SetFilePointerEx(out->hFile, outputSectionsSize, NULL,
				FILE_BEGIN)) {
			CloseMappedFile(out);
			return FALSE;
		}
		if (!WriteFile(out->hFile, in->mapping + peInfo->sizeOfImageInFile,
				inputFileSize.QuadPart - peInfo->sizeOfImageInFile,
				&bytesWritten, NULL)) {
			CloseMappedFile(out);
			return FALSE;
		}
		if (bytesWritten
				!= inputFileSize.QuadPart - peInfo->sizeOfImageInFile) {
			CloseMappedFile(out);
			return FALSE;
		}
	}
	out->mapSize = peInfo->sizeOfImageInFile + sizeIncrease;
	out->hMapping = CreateFileMapping(out->hFile, NULL, PAGE_READWRITE, 0,
			out->mapSize, NULL);
	if (out->hMapping == INVALID_HANDLE_VALUE) {
		CloseMappedFile(out);
		return FALSE;
	}
	out->mapping = MapViewOfFile(out->hMapping, FILE_MAP_WRITE, 0, 0,
			out->mapSize);
	if (out->mapping == NULL) {
		CloseMappedFile(out);
		return FALSE;
	}
	return TRUE;
}

void CopySections(mapped_file* out, pe_info* peInfo) {
	int i;
	int fileStart;
	int fileSize;
	int outputStart;
	for (i = 0; i < peInfo->numberOfSections; i++) {
		fileStart = peInfo->sections[i].PointerToRawData;
		fileSize = AlignUp(peInfo->sections[i].SizeOfRawData,
				peInfo->fileAlignment);
		outputStart = fileStart + peInfo->fileOffset;
		memcpy(out->mapping + outputStart, peInfo->mapping + fileStart,
				fileSize);
	}
	memcpy(out->mapping + peInfo->sizeOfImageInFile + peInfo->fileOffset,
			peInfo->newSection.data, peInfo->newSection.size);
}

void UpdateHeader32(mapped_file* out, pe_info* peInfo,
		PIMAGE_SECTION_HEADER newSection) {
	PIMAGE_NT_HEADERS32 peHeader;
	int optionalHeaderOffset;
	int i;
	PIMAGE_SECTION_HEADER newSections;
	peHeader = (PIMAGE_NT_HEADERS32) (out->mapping + peInfo->headerOffset);
	peHeader->FileHeader.NumberOfSections++;
	optionalHeaderOffset = (((char*) &peHeader->OptionalHeader) - out->mapping);
	peHeader->FileHeader.SizeOfOptionalHeader = peInfo->sectionStartOffset
			- optionalHeaderOffset;
	peHeader->OptionalHeader.SizeOfHeaders =
			max(peInfo->sectionEndOffset + sizeof(IMAGE_SECTION_HEADER),
					peHeader->OptionalHeader.SizeOfHeaders);
	peHeader->OptionalHeader.SizeOfImage += AlignUp(peInfo->newSection.size,
			peInfo->sectionAlignment);
	if (peInfo->rvaOffset > 0) {
		newSections = IMAGE_FIRST_SECTION(peHeader);
		for (i = 0; i < peHeader->FileHeader.NumberOfSections; i++) {
			newSections[i].VirtualAddress += peInfo->rvaOffset;
		}
	}
}

void UpdateHeader64(mapped_file* out, pe_info* peInfo,
		PIMAGE_SECTION_HEADER newSection) {
	PIMAGE_NT_HEADERS64 peHeader;
	int optionalHeaderOffset;
	int i;
	PIMAGE_SECTION_HEADER newSections;
	peHeader = (PIMAGE_NT_HEADERS64) (out->mapping + peInfo->headerOffset);
	peHeader->FileHeader.NumberOfSections++;
	optionalHeaderOffset = (((char*) &peHeader->OptionalHeader) - out->mapping);
	peHeader->FileHeader.SizeOfOptionalHeader = peInfo->sectionStartOffset
			- optionalHeaderOffset;
	peHeader->OptionalHeader.SizeOfHeaders =
			max(peInfo->sectionEndOffset + sizeof(IMAGE_SECTION_HEADER),
					peHeader->OptionalHeader.SizeOfHeaders);
	peHeader->OptionalHeader.SizeOfImage += AlignUp(peInfo->newSection.size,
			peInfo->sectionAlignment);
	if (peInfo->rvaOffset > 0) {
		newSections = IMAGE_FIRST_SECTION(peHeader);
		for (i = 0; i < peHeader->FileHeader.NumberOfSections; i++) {
			newSections[i].VirtualAddress += peInfo->rvaOffset;
		}
	}
}

void CopyHeader(mapped_file* out, pe_info* peInfo) {
	PIMAGE_SECTION_HEADER newSection;
	IMAGE_DATA_DIRECTORY* newDirectory;
	memcpy(out->mapping, peInfo->mapping, peInfo->sectionEndOffset);
	newSection = (PIMAGE_SECTION_HEADER) (out->mapping
			+ peInfo->sectionEndOffset);
	newSection->VirtualAddress = peInfo->newSectionRVA;
	newSection->PointerToRawData = peInfo->sizeOfImageInFile;
	newSection->SizeOfRawData = peInfo->newSection.size;
	newSection->Misc.VirtualSize = AlignUp(peInfo->newSection.size,
			peInfo->sectionAlignment);
	newSection->PointerToLinenumbers = 0;
	newSection->PointerToRelocations = 0;
	newSection->NumberOfLinenumbers = 0;
	newSection->NumberOfRelocations = 0;
	strcpy(newSection->Name, ".ihook");
	newSection->Characteristics = IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE;
	if (peInfo->bits == 32) {
		UpdateHeader32(out, peInfo, newSection);
	} else {
		UpdateHeader64(out, peInfo, newSection);
	}
	newDirectory = (IMAGE_DATA_DIRECTORY*) (out->mapping
			+ (((char*) peInfo->directory) - peInfo->mapping));
	newDirectory[IMAGE_DIRECTORY_ENTRY_BOUND_IMPORT].Size = 0;
	newDirectory[IMAGE_DIRECTORY_ENTRY_BOUND_IMPORT].VirtualAddress = 0;
	newDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress =
			peInfo->importTablesRVA;
	newDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].Size =
			(peInfo->newNumberOfImportTables + 1)
					* sizeof(IMAGE_IMPORT_DESCRIPTOR);
}

int _tmain(int argc, _TCHAR **argv) {
	mapped_file in;
	mapped_file out;
	pe_info peInfo;
	if (argc <= 4) {
		OutputMessage(usage);
		return 0;
	}
	if (!OpenMappedFile(&in, argv[1])) {
		CloseMappedFile(&in);
		return 1;
	}
	if (!ParseInputFileHeader(&in, &peInfo)) {
		CloseMappedFile(&in);
		return 1;
	}
	if (!AddImportTables(&in, &peInfo, &argv[3], argc - 3)) {
		CloseMappedFile(&in);
		return 1;
	}
	if (!CreateOutputFile(&in, &out, &peInfo, argv[2])) {
		CloseMappedFile(&in);
		return 1;
	}
	CopySections(&out, &peInfo);
	CopyHeader(&out, &peInfo);
	CloseMappedFile(&out);
	CloseMappedFile(&in);
	return 0;
}

